"""
调度器模块

负责协调下载器、解析器和存储模块的工作，管理任务队列
"""

import os
import time
import queue
import logging
import threading
from urllib.parse import urlparse
from collections import defaultdict

logger = logging.getLogger('scheduler')

class Scheduler:
    """调度器类"""
    
    def __init__(self, config, downloader, parser_manager, storage):
        """
        初始化调度器
        
        参数:
            config: 配置对象
            downloader: 下载器对象
            parser_manager: 解析器管理器对象
            storage: 存储模块对象
        """
        self.config = config
        self.downloader = downloader
        self.parser_manager = parser_manager
        self.storage = storage
        
        # 从配置中读取设置
        self.max_depth = int(config.get('DEFAULT', 'max_depth', fallback=1))
        self.same_domain_only = config.getboolean('DEFAULT', 'same_domain_only', fallback=True)
        
        # 创建任务队列
        self.task_queue = queue.Queue()
        
        # 创建已访问URL集合
        self.visited_urls = set()
        
        # 创建已加入队列但尚未访问的URL集合
        self.pending_urls = set()
        
        # 创建URL深度映射
        self.url_depths = {}
        
        # 创建统计信息
        self.stats = {
            'total_tasks': 0,
            'successful_tasks': 0,
            'failed_tasks': 0,
            'visited_urls': 0
        }
        
        # 创建运行标志
        self.running = False
        
        # 创建工作线程
        self.worker_thread = None
        
        # 创建完成事件
        self.done_event = threading.Event()
        
        # 导入链接提取器和翻页处理器
        from crawler.link_extractor import LinkExtractorManager
        from crawler.paginator import PaginatorManager
        
        # 创建链接提取器和翻页处理器
        self.link_extractor = LinkExtractorManager(config)
        self.paginator = PaginatorManager(config)
        
        logger.info("调度器初始化完成")
    
    def add_task(self, url, depth=0):
        """
        添加爬取任务
        
        参数:
            url: 要爬取的URL
            depth: 当前爬取深度
        
        返回:
            布尔值，表示是否成功添加
        """
        # 检查URL是否已访问或已在队列中
        if url in self.visited_urls or url in self.pending_urls:
            logger.debug(f"URL已访问或已在队列中，跳过: {url}")
            return False
        
        # 检查是否超过最大深度
        if depth > self.max_depth:
            logger.debug(f"超过最大深度({self.max_depth})，跳过: {url}")
            return False
        
        # 检查是否仅爬取同域名
        if self.same_domain_only and depth > 0:
            # 获取初始域名（深度为0的URL的域名）
            initial_domains = set()
            for u, d in self.url_depths.items():
                if d == 0:
                    initial_domains.add(urlparse(u).netloc)
            
            # 如果没有初始域名，将当前URL作为初始域名
            if not initial_domains:
                initial_domains.add(urlparse(url).netloc)
            
            # 检查当前URL是否属于初始域名
            current_domain = urlparse(url).netloc
            if not any(current_domain.endswith(domain) for domain in initial_domains):
                logger.debug(f"非同域名URL，跳过: {url}")
                return False
        
        # 添加到任务队列
        self.task_queue.put((url, depth))
        self.url_depths[url] = depth
        self.pending_urls.add(url)  # 标记为待处理
        self.stats['total_tasks'] += 1
        
        logger.info(f"添加任务: {url} (深度: {depth})")
        return True
    
    def add_tasks_from_file(self, file_path):
        """
        从文件中批量添加任务
        
        参数:
            file_path: 包含URL列表的文件路径
        
        返回:
            成功添加的任务数
        """
        count = 0
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    url = line.strip()
                    if url and not url.startswith('#'):
                        if self.add_task(url):
                            count += 1
        except Exception as e:
            logger.error(f"从文件添加任务时出错: {e}")
        
        return count
    
    def start(self):
        """启动调度器"""
        if self.running:
            logger.warning("调度器已在运行")
            return
        
        self.running = True
        self.done_event.clear()
        
        # 创建并启动工作线程
        self.worker_thread = threading.Thread(target=self._worker_loop)
        self.worker_thread.daemon = True
        self.worker_thread.start()
        
        logger.info("调度器已启动")
    
    def stop(self):
        """停止调度器"""
        if not self.running:
            logger.warning("调度器未在运行")
            return
        
        self.running = False
        
        # 等待工作线程结束
        if self.worker_thread and self.worker_thread.is_alive():
            self.worker_thread.join(timeout=5)
        
        logger.info("调度器已停止")
    
    def wait_until_done(self, timeout=None):
        """
        等待所有任务完成
        
        参数:
            timeout: 超时时间(秒)，None表示无限等待
        
        返回:
            布尔值，表示是否所有任务都已完成
        """
        return self.done_event.wait(timeout)
    
    def get_stats(self):
        """
        获取统计信息
        
        返回:
            统计信息字典
        """
        return self.stats
    
    def _worker_loop(self):
        """工作线程循环"""
        while self.running:
            try:
                # 检查任务队列是否为空
                if self.task_queue.empty():
                    logger.info("任务队列为空，工作线程退出")
                    self.done_event.set()
                    break
                
                # 从队列中获取任务
                url, depth = self.task_queue.get(block=False)
                
                # 从待处理集合中移除
                self.pending_urls.discard(url)
                
                # 标记URL为已访问
                self.visited_urls.add(url)
                self.stats['visited_urls'] += 1
                
                logger.info(f"处理任务: {url} (深度: {depth})")
                
                # 下载页面
                html = self.downloader.download(url)
                if not html:
                    logger.warning(f"下载失败: {url}")
                    self.stats['failed_tasks'] += 1
                    continue
                
                # 解析页面
                data = self.parser_manager.parse(html, url)
                if not data:
                    logger.warning(f"解析失败: {url}")
                    self.stats['failed_tasks'] += 1
                    continue
                
                # 存储数据
                self.storage.save(data, url)
                self.stats['successful_tasks'] += 1
                
                # 提取链接并添加到任务队列
                self._extract_and_follow_links(html, url, depth)
                
                # 处理翻页
                self._handle_pagination(html, url, depth)
                
            except queue.Empty:
                # 队列为空，等待一会儿再试
                time.sleep(0.1)
            except Exception as e:
                logger.error(f"处理任务时出错: {e}")
                self.stats['failed_tasks'] += 1
    
    def _extract_and_follow_links(self, html, url, depth):
        """
        提取链接并添加到任务队列
        
        参数:
            html: HTML内容
            url: 当前页面URL
            depth: 当前爬取深度
        """
        # 检查是否已达到最大深度
        if depth >= self.max_depth:
            logger.debug(f"已达到最大深度({self.max_depth})，不再提取链接")
            return
        
        # 提取链接
        links = self.link_extractor.extract_links(html, url)
        
        # 添加链接到任务队列
        added_count = 0
        for link in links:
            if self.add_task(link, depth + 1):
                added_count += 1
        
        logger.info(f"从 {url} 提取并添加了 {added_count} 个链接任务")
    
    def _handle_pagination(self, html, url, depth):
        """
        处理翻页
        
        参数:
            html: HTML内容
            url: 当前页面URL
            depth: 当前爬取深度
        """
        # 获取下一页URL
        next_page = self.paginator.get_next_page(html, url)
        
        # 如果有下一页，添加到任务队列
        if next_page:
            # 翻页任务使用相同的深度，不增加深度计数
            if self.add_task(next_page, depth):
                logger.info(f"添加翻页任务: {next_page} (深度: {depth})")
